package com.pig4cloud.pigx.knowledge.support.handler.rag;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.NumberUtil;
import cn.hutool.core.util.StrUtil;
import com.pig4cloud.pigx.knowledge.dto.AiMessageResultDTO;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.entity.AiDatasetEntity;
import com.pig4cloud.pigx.knowledge.entity.AiDocumentEntity;
import com.pig4cloud.pigx.knowledge.entity.AiSliceEntity;
import com.pig4cloud.pigx.knowledge.service.AiDocumentService;
import com.pig4cloud.pigx.knowledge.service.AiSliceService;
import com.pig4cloud.pigx.knowledge.service.EmbeddingStoreService;
import com.pig4cloud.pigx.knowledge.support.constant.DocumentTypeEnums;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import lombok.SneakyThrows;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * 问题 搜索向量处理器
 *
 * @author lengleng
 * @date 2024/7/4
 */
@Service
public class Q2AVectorRagChatHandler extends AbstractRagChatHandler {

	private final EmbeddingStoreService embeddingStoreService;

	private final AiDocumentService aiDocumentService;

	private final AiSliceService aiSliceService;

	public Q2AVectorRagChatHandler(ModelProvider modelProvider, AiDocumentService aiDocumentService,
			AiSliceService aiSliceService, EmbeddingStoreService embeddingStoreService) {
		super(modelProvider);
		this.aiDocumentService = aiDocumentService;
		this.aiSliceService = aiSliceService;
		this.embeddingStoreService = embeddingStoreService;
	}

	public Flux<AiMessageResultDTO> process(Embedding embeddedList, AiDatasetEntity dataset,
			ChatMessageDTO chatMessageDTO) {
		double minScore = NumberUtil.div(Double.parseDouble(dataset.getScore().toString()), Double.parseDouble("100"),
				2);

		EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
			.queryEmbedding(embeddedList)
			.maxResults(dataset.getTopK())
			.filter(metadataKey(AiDocumentEntity.Fields.datasetId).isEqualTo(dataset.getId().toString())
				.and(metadataKey(DocumentTypeEnums.Fields.type).isEqualTo(DocumentTypeEnums.ANSWER.getType())))
			.minScore(minScore)
			.build();

		EmbeddingSearchResult<TextSegment> searchResult = embeddingStoreService
			.embeddingStore(dataset.getCollectionName())
			.search(embeddingSearchRequest);
		List<EmbeddingMatch<TextSegment>> embeddingMatchList = searchResult.matches();

		// 未匹配
		if (CollUtil.isEmpty(embeddingMatchList)) {
			return Flux.just(new AiMessageResultDTO(dataset.getEmptyDesc()));
		}

		// 更新命中次数
		List<String> embeddingIdList = embeddingMatchList.stream().map(EmbeddingMatch::embeddingId).toList();
		aiSliceService.updateHitCount(embeddingIdList);
		// 对向量结果进行总结
		Flux<AiMessageResultDTO> aiMessageResultDTOFlux = summaryResult(dataset, chatMessageDTO,
				embeddingMatchList.stream().map(EmbeddingMatch::embedded).map(TextSegment::text).toList())
			.cache();

		// 修改 map 逻辑在最后拼接一下参考资料
		AiMessageResultDTO aiMessageResultDTO = new AiMessageResultDTO();
		aiMessageResultDTO.setMessage(StrUtil.EMPTY);
		List<AiMessageResultDTO.ExtLink> extLinks = buildExtMessage(embeddingMatchList);
		aiMessageResultDTO.setExtLinks(extLinks);
		return aiMessageResultDTOFlux.concatWithValues(aiMessageResultDTO);
	}

	/**
	 * 构建扩展信息 （参考资料 tag）
	 * @return String
	 */
	@SneakyThrows
	private List<AiMessageResultDTO.ExtLink> buildExtMessage(List<EmbeddingMatch<TextSegment>> embeddingMatchList) {
		List<AiMessageResultDTO.ExtLink> extLinks = new ArrayList<>();
		for (EmbeddingMatch<TextSegment> textSegmentEmbeddingMatch : embeddingMatchList) {

			String documentId = textSegmentEmbeddingMatch.embedded()
				.metadata()
				.getString(AiSliceEntity.Fields.documentId);
			if (Objects.isNull(documentId)) {
				continue;
			}

			Float distance = textSegmentEmbeddingMatch.score().floatValue();
			AiDocumentEntity aiDocumentEntity = aiDocumentService.getById(Long.parseLong(documentId));
			extLinks.add(new AiMessageResultDTO.ExtLink(aiDocumentEntity.getName(),
					String.format("/admin/sys-file/oss/file?fileName=%s", aiDocumentEntity.getFileUrl()), distance));
		}

		return extLinks.stream().distinct().toList();
	}

}
